/* Copyright 2022 The OpenXLA Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef XLA_PJRT_PJRT_FUTURE_H_
#define XLA_PJRT_PJRT_FUTURE_H_

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <functional>
#include <memory>
#include <type_traits>
#include <utility>

#include "absl/base/no_destructor.h"
#include "absl/base/optimization.h"
#include "absl/meta/type_traits.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "absl/utility/utility.h"
#include "xla/tsl/concurrency/async_value.h"
#include "xla/tsl/concurrency/async_value_ref.h"
#include "xla/tsl/concurrency/ref_count.h"
#include "xla/tsl/platform/logging.h"

namespace xla {

template <class T = void>
class PjRtFuture;

namespace internal {
template <class T, bool is_move_only>
class PjRtFutureBase;
}

// Returns a `PjRtFuture` that will be successful if all `futures` complete
// successfully, or return a first encountered error.
PjRtFuture<> JoinFutures(absl::Span<const PjRtFuture<>> futures);

// An RAII event that a caller can use to tell the PjRtClient about asynchronous
// actions outside PjRt.
//
// A ScopedAsyncTrackingEvent can be generated by the caller by calling a method
// on PjRtDevice, and the creation of a ScopedAsyncTrackingEvent tells the
// PjRtClient that the client is creating some outstanding asynchronous work
// that depends on activities happening on the PjRtDevice.
//
// The caller can indicate that a ScopedAsyncTrackingEvent event cannot complete
// until after some PjRtFuture becomes ready, by calling
// future.AssertHappensBefore(event).
//
// The caller indicates that the work tracked by the ScopedAsyncTrackingEvent
// has completed by letting the event go out of scope.
//
// ScopedAsyncTrackingEvent is used by some PjRtClient implementations to
// monitor system-wide dependencies.
class ScopedAsyncTrackingEvent {
 public:
  virtual ~ScopedAsyncTrackingEvent() = default;

 private:
  template <class T, bool is_move_only>
  friend class internal::PjRtFutureBase;

  // Indicates that the ScopedAsyncTrackingEvent won't complete until dependency
  // becomes available. Called only by PjRtFuture.
  virtual void AddDependency(tsl::RCReference<tsl::AsyncValue> dependency) = 0;
};

// Helpers for using PjRtFutures.
class PjRtFutureHelpers {
 public:
  // Keys that are returned by an implementation-specific handler when a client
  // starts to block on a promise.
  //
  // For now, contains a single UID that can be used to identify a TraceMe, but
  // made extensible to allow support for other profilers such as endoscope.
  struct ProfilingKeys {
    uint64_t traceme_context_id = -1;
  };

  // Signature of handler called by the PjRtFuture class before it starts to
  // block a thread.
  using OnBlockStartFn = std::function<ProfilingKeys()>;

  // Signature of handler called by the PjRtFuture class after it finishes
  // blocking a thread.
  using OnBlockEndFn = std::function<void(ProfilingKeys)>;

  // Returns a PjRtFuture<T> with optionally updated profiling handlers. If
  // profiling handlers are not provided, the original ones will be used.
  template <typename T>
  static PjRtFuture<T> WithProfiling(PjRtFuture<T> future,
                                     OnBlockStartFn on_block_start = nullptr,
                                     OnBlockEndFn on_block_end = nullptr) {
    return PjRtFuture<T>(std::move(future.promise_),
                         on_block_start ? std::move(on_block_start)
                                        : std::move(future.on_block_start_),
                         on_block_end ? std::move(on_block_end)
                                      : std::move(future.on_block_end_));
  }
};

namespace internal {

// Detects absl::StatusOr<T> specializations to disable them for PjRtFuture<T>.
template <typename T>
struct IsStatusOr : public std::false_type {};
template <typename T>
struct IsStatusOr<absl::StatusOr<T>> : public std::true_type {};

// A base class to conditionally disable copy constructor and assignment for a
// PjRtFuture<T> (by default we always disable copy constructor when `T` is not
// copyable), which makes PjRtFuture<T> an `std::unique_ptr`-like container for
// move-only types.
template <bool is_move_only>
class PjRtFutureMoveControl;

template <>
class PjRtFutureMoveControl</*is_move_only=*/true> {
 protected:
  PjRtFutureMoveControl() = default;

  PjRtFutureMoveControl(const PjRtFutureMoveControl&) = delete;
  PjRtFutureMoveControl& operator=(const PjRtFutureMoveControl&) = delete;

  PjRtFutureMoveControl(PjRtFutureMoveControl&&) = default;
  PjRtFutureMoveControl& operator=(PjRtFutureMoveControl&&) = default;
};

template <>
class PjRtFutureMoveControl</*is_move_only=*/false> {
 protected:
  PjRtFutureMoveControl() = default;

  PjRtFutureMoveControl(const PjRtFutureMoveControl&) = default;
  PjRtFutureMoveControl& operator=(const PjRtFutureMoveControl&) = default;

  PjRtFutureMoveControl(PjRtFutureMoveControl&&) = default;
  PjRtFutureMoveControl& operator=(PjRtFutureMoveControl&&) = default;
};

// A base class for a stateful future PjRtFuture<T> and a stateless future
// PjRtFuture<>. If `is_move_only` is true, PjRtFuture derived from this class
// acts as a move-only type and the value can be passed to the caller only using
// move assignment (applied to Await and OnReady APIs).
template <typename T, bool is_move_only = !std::is_copy_constructible_v<T>>
class PjRtFutureBase : public PjRtFutureMoveControl<is_move_only> {
 protected:
  // A protected constructor that hides AsyncValueRef implementation detail
  // from the end users of PjRtFuture and Promise. Must not be made public!
  PjRtFutureBase(tsl::AsyncValueRef<T> promise,
                 PjRtFutureHelpers::OnBlockStartFn on_block_start,
                 PjRtFutureHelpers::OnBlockEndFn on_block_end)
      : promise_(std::move(promise)),
        on_block_start_(std::move(on_block_start)),
        on_block_end_(std::move(on_block_end)) {}

 public:
  PjRtFutureBase() = default;

  // Constructor for an already-available PjRtFuture.
  //
  // Typically used to eagerly return error values when async work will not
  // be enqueued, e.g., due to invalid arguments.
  explicit PjRtFutureBase(
      T t, PjRtFutureHelpers::OnBlockStartFn on_block_start = nullptr,
      PjRtFutureHelpers::OnBlockEndFn on_block_end = nullptr)
      : PjRtFutureBase(tsl::MakeAvailableAsyncValueRef<T>(std::move(t)),
                       std::move(on_block_start), std::move(on_block_end)) {}

  bool IsValid() const { return promise_ != nullptr; }

  // Two functions exist to know whether the future is ready, to accommodate
  // the fact some backends (e.g. distributed ones) could take a non-trivial
  // time to check the state of a future.
  //
  // `IsReady()` is guaranteed to return true if the future became ready
  // before `IsReady()` was called. `IsReady()` will return immediately if a
  // call to `Await()` has already returned, or any callback passed to
  // `OnReady` has already been triggered. Otherwise IsReady() may block for
  // the duration of a network message on some backends.
  bool IsReady() const {
    CHECK(IsValid());
    return promise_.IsAvailable();
  }
  // `IsKnownReady()` is guaranteed to return immediately. `IsKnownReady()` will
  // always return true if a call to `Await()` has already returned, or any
  // callback passed to `OnReady` has already been triggered. Otherwise,
  // `IsKnownReady()` may return false in some cases in which the future was
  // ready before `IsKnownReady()` was called.
  bool IsKnownReady() const {
    CHECK(IsValid());
    return promise_.IsAvailable();
  }

  // Indicates that event will not complete until after this becomes ready.
  //
  // May safely be called with event==nullptr in which case AssertHappensBefore
  // has no effect.
  void AssertHappensBefore(ScopedAsyncTrackingEvent* event) {
    CHECK(IsValid());
    if (event) {
      event->AddDependency(promise_.CopyRCRef());
    }
  }

  explicit operator bool() const { return static_cast<bool>(promise_); }

 protected:
  static constexpr bool IsMoveOnly() { return is_move_only; }

  // PjRtFuture<T>::Promise provides a facility to store a value or an error
  // that is later acquired asynchronously via a PjRtFuture<T> constructed from
  // the promise object. Note that the promise object is meant to be used only
  // once (set value or error).
  class Promise {
   public:
    Promise() = default;

    Promise(Promise&& other) = default;
    Promise& operator=(Promise&& other) = default;

    Promise(const Promise& other) = default;
    Promise& operator=(const Promise& other) = default;

    explicit operator bool() const { return static_cast<bool>(promise_); }

    // Returns if this promise is the unique reference to the underlying value.
    // That is, this method returns true only if all of the following conditions
    // are satisfied:
    //
    // - The promise is the only reference to the underlying value, i.e., there
    //   are no other promises or futures associated with this value.
    // - There are no OnReady callbacks registered to this promise.
    //
    // This may be used by the caller of `Set()` to short-circuit the work to
    // fulfill the promise if no one will ever consume the value. Even in that
    // case, consider fulfilling the promise with an error (e.g., `CANCELLED`)
    // instead of dropping the promise without fulfilling it in order to make
    // debugging easier. Also, be aware that the current promise may still be
    // used to mint a future.
    bool IsUniqueReference() const {
      return async_value()->IsUnique() && !async_value()->HasWaiter();
    }

   protected:
    explicit Promise(tsl::AsyncValueRef<T> promise)
        : promise_(std::move(promise)) {}

    template <typename... Args>
    void emplace(Args&&... args) const {
      DCHECK(promise_) << "Promise must wrap an async value";
      promise_.template emplace<Args...>(std::forward<Args>(args)...);
    }

    // Takes a reference to the underlying AsyncValueRef container.
    tsl::AsyncValueRef<T> ref() const { return promise_; }

    // Returns a pointer to the underlying AsyncValue that can be used to
    // track completion of a promise. It is undefined behavior to access the
    // value stored in the AsyncValue.
    tsl::AsyncValue* async_value() const { return promise_.GetAsyncValue(); }

#ifndef NDEBUG
    int64_t AddFuture() const { return num_futures_->fetch_add(1); }
#endif

   private:
    tsl::AsyncValueRef<T> promise_;

#ifndef NDEBUG
    // In debug builds we track the number of futures created from a promise to
    // detect when a promise for a move-only type can be accidentally shared by
    // multiple futures. We wrap the counter into shared pointer because promise
    // for a move-only future is still copyable, but only one future can be
    // created from all the copies.
    std::shared_ptr<std::atomic<int64_t>> num_futures_ =
        std::make_shared<std::atomic<int64_t>>(0);
#endif
  };

  class ProfilingCleanup {
   public:
    ProfilingCleanup(const PjRtFutureBase* parent,
                     PjRtFutureHelpers::ProfilingKeys keys)
        : parent_(parent), keys_(std::move(keys)) {}
    ~ProfilingCleanup() {
      if (parent_ && parent_->on_block_end_) {
        parent_->on_block_end_(std::move(keys_));
      }
    }
    ProfilingCleanup(const ProfilingCleanup& other) = delete;
    ProfilingCleanup(ProfilingCleanup&& other) = delete;

   private:
    const PjRtFutureBase* parent_;
    PjRtFutureHelpers::ProfilingKeys keys_;
  };

  ProfilingCleanup OnBlockStartScope() const {
    return ProfilingCleanup(this, on_block_start_
                                      ? on_block_start_()
                                      : PjRtFutureHelpers::ProfilingKeys());
  }

  // Calls block_until_ready_fn to wait until the underlying AsyncValue is
  // concrete. block_until_ready_fn should be equivalent to
  // tsl::BlockUntilReady.
  template <typename Fn>
  void BlockUntilReady(Fn&& block_until_ready_fn) const {
    CHECK(IsValid());
    if (!promise_.IsAvailable()) {
      ProfilingCleanup scope = OnBlockStartScope();
      block_until_ready_fn(promise_.GetAsyncValue());
    }
    DCHECK(promise_.IsConcrete());
  }

  // Blocks the calling thread until the future is ready, then returns the
  // final value.
  const T& Await() const& {
    BlockUntilReady(
        static_cast<void (*)(tsl::AsyncValue*)>(tsl::BlockUntilReady));
    return *promise_;
  }

  // Blocks the calling thread until the future is ready, then returns the
  // final value.
  std::conditional_t<is_move_only, T, const T&> Await() && {
    BlockUntilReady(
        static_cast<void (*)(tsl::AsyncValue*)>(tsl::BlockUntilReady));

    if constexpr (is_move_only) {
      return std::move(*promise_);
    } else {
      // We can't move from the promise to the caller because for copyable
      // futures we can have multiple copies of the PjRtFuture sharing the
      // same underlying promise object.
      return *promise_;
    }
  }

  // Returns a PjRtFuture<> that becomes ready when *this is ready. If *this
  // completes with an error, the returned future will also be an error.
  //
  // This function defined out of line as it requires PjRtFuture<> definition.
  PjRtFuture<> GetReadyFuture() const;

  // Registers callback to be called once the promise is ready, with the final
  // value.
  //
  // callback may be called on an internal system thread or the calling thread.
  // The client should avoid any potentially re-entrant API calls within the
  // callback, for example by using the callback to enqueue work on a
  // client-owned threadpool.
  template <typename F,
            std::enable_if_t<!is_move_only &&
                             std::is_invocable_v<F, const T&>>* = nullptr>
  void OnReady(F&& f) const& {
    CHECK(IsValid());
    promise_.AndThen(
        [promise = promise_.AsPtr(), f = std::forward<F>(f)]() mutable {
          DCHECK(promise.IsConcrete());
          f(*promise);
        });
  }

  // Registers callback to be called once the promise is ready, with the final
  // value.
  //
  // callback may be called on an internal system thread or the calling thread.
  // The client should avoid any potentially re-entrant API calls within the
  // callback, for example by using the callback to enqueue work on a
  // client-owned threadpool.
  template <typename F,
            std::enable_if_t<is_move_only ? std::is_invocable_v<F, T>
                                          : std::is_invocable_v<F, const T&>>* =
                nullptr>
  void OnReady(F&& f) && {
    CHECK(IsValid());
    promise_.AndThen(
        [promise = promise_.AsPtr(), f = std::forward<F>(f)]() mutable {
          DCHECK(promise.IsConcrete());
          if constexpr (is_move_only) {
            f(std::move(*promise));
          } else {
            // We can't move from the promise to the caller because for copyable
            // futures we can have multiple copies of the PjRtFuture sharing the
            // same underlying promise object.
            f(*promise);
          }
        });
  }

 protected:
  // Returns a placeholder error that can be used when short-circuiting promises
  // with no other references.
  static absl::Status AbortedError() {
    return absl::AbortedError(
        "Fulfilling the promise with an aborted error since the value is no "
        "longer referenced by any futures or OnReady callbacks; if this error "
        "is exposed to any future, that indicates a bug");
  }

 private:
  friend class xla::PjRtFutureHelpers;

  tsl::AsyncValueRef<T> promise_;

  // Function that is called before a thread starts blocking on the promise.
  PjRtFutureHelpers::OnBlockStartFn on_block_start_;
  // Function that is called after a thread finishes blocking on the promise.
  PjRtFutureHelpers::OnBlockEndFn on_block_end_;
};

}  // namespace internal

// PjRtFuture<T> is a simple future that is returned by PjRt APIs that
// enqueue asynchronous work, reporting a value of type T when the work is
// complete.
//
// PjRtFuture can be used by the client to wait for work to complete, either via
// a blocking call or a callback.
//
// The implementation wraps a tsl::AsyncValueRef<T>, but we prefer to
// encapsulate the AVR rather than returning it directly for three reasons.
//
// First, in contrast to AsyncValueRef which has a smart-pointer semantics,
// future has more of a value semantics, i.e. future of a move-only type also
// is a move-only type. You can think of a move-only future as a box to pass a
// value of type T between asynchronous producer/consumer: you can open the box
// once to put the value into it and you can open the box only once to take the
// value out of it. For copyable types PjRtFuture<T> is a copyable type,
// although all copies share the same underlying value.
//
// Second, we want to retain portability in case a future implementation moves
// away from AsyncValueRef ---- we don't want clients to call arbitrary
// AsyncValueRef APIs.
//
// Third, we want to export different semantics, for example we support
// integration between blocking and profiling (e.g., TraceMe).
template <class T>
class PjRtFuture : public internal::PjRtFutureBase<absl::StatusOr<T>> {
  using Base = internal::PjRtFutureBase<absl::StatusOr<T>>;

  static constexpr bool is_move_only = Base::IsMoveOnly();  // NOLINT

  template <typename U>
  static constexpr bool is_status_or =  // NOLINT
      tsl::internal::is_status_or_v<U>;

  static_assert(!std::is_same_v<T, absl::Status>,
                "Use PjRtFuture<> specialization for stateless futures");

  static_assert(
      !internal::IsStatusOr<T>::value,
      "PjRtFuture<T> already has an implicit absl::StatusOr<T> semantics");

 public:
  class Promise : public Base::Promise {
   public:
    using Base::Promise::Promise;

    // Sets the value of the promise. Must be called at most once.
    //
    // After Set is called, value will be delivered to waiters on the PjRtFuture
    // constructed from a promise, via blocking or callbacks.
    void Set(absl::StatusOr<T> value) {
      Base::Promise::emplace(std::move(value));
    }

   private:
    template <typename>
    friend class PjRtFuture;
  };

  // This is a temporary class to support migration from CreatePromise() to
  // MakePromise() and an end goal of making Promise move-only type.
  class MoveOnlyPromise : public Promise {
   public:
    using Promise::Promise;
    using Promise::Set;

    MoveOnlyPromise(MoveOnlyPromise&&) = default;
    MoveOnlyPromise& operator=(MoveOnlyPromise&&) = default;

    // A helper function to convert move-only Promise to shared_ptr, which is
    // useful when the promise has to be captured by a std::function.
    std::shared_ptr<MoveOnlyPromise> ToShared() && {
      return std::make_shared<MoveOnlyPromise>(std::move(*this));
    }
  };

  // Returns a Promise that can be used to construct a PjRtFuture, and then Set
  // later.
  static Promise CreatePromise() {
    return Promise(tsl::MakeUnconstructedAsyncValueRef<absl::StatusOr<T>>());
  }

  // Returns a pair of connected Promise and PjRtFuture<T>. Setting the returned
  // promise will fulfill the connected future.
  static std::pair<MoveOnlyPromise, PjRtFuture<T>> MakePromise(
      PjRtFutureHelpers::OnBlockStartFn on_block_start = nullptr,
      PjRtFutureHelpers::OnBlockEndFn on_block_end = nullptr) {
    MoveOnlyPromise promise(
        tsl::MakeUnconstructedAsyncValueRef<absl::StatusOr<T>>());
    PjRtFuture<T> future(promise, std::move(on_block_start),
                         std::move(on_block_end));
    return std::make_pair(std::move(promise), std::move(future));
  }

  // Bring PjRtFutureBase constructors in scope.
  using Base::Base;

  // Constructor for unavailable future that will be fulfilled later via the
  // promise object.
  //
  // - on_block_start is called before Await starts to block.
  //  - on_block_end is called after Await finishes blocking.
  explicit PjRtFuture(
      const Promise& promise,
      PjRtFutureHelpers::OnBlockStartFn on_block_start = nullptr,
      PjRtFutureHelpers::OnBlockEndFn on_block_end = nullptr)
      : Base(promise.ref(), std::move(on_block_start),
             std::move(on_block_end)) {
#ifndef NDEBUG
    if constexpr (is_move_only) {
      DCHECK_EQ(promise.AddFuture(), 0)
          << "Move-only PjRtFuture cannot share a promise object";
    }
#endif
  }

  using Base::Await;
  using Base::GetReadyFuture;
  using Base::OnReady;

  // Returns an PjRtFuture<R> that is constructed from the result of invoking
  // functor `f` with *this value. If *this completes with an error, returned
  // future will also be an error.
  //
  // Sample usage:
  //
  // future.Map<R>([](const T& value) -> U {
  //   return U(value); // R must be constructible from U
  // })
  //
  template <typename R, typename F,
            typename U = std::invoke_result_t<F, const T&>,
            std::enable_if_t<!is_move_only && std::is_constructible_v<R, U>>* =
                nullptr>
  PjRtFuture<R> Map(F&& f) const& {
    auto [promise, future] = PjRtFuture<R>::MakePromise();

    using Value = const absl::StatusOr<T>&;
    OnReady([promise = std::move(promise),
             f = std::forward<F>(f)](Value value) mutable {
      if (ABSL_PREDICT_FALSE(promise.IsUniqueReference())) {
        promise.Set(Base::AbortedError());
      } else if (ABSL_PREDICT_TRUE(value.ok())) {
        promise.emplace(absl::in_place_t{}, f(*value));
      } else {
        promise.Set(value.status());
      }
    });

    return std::move(future);
  }

  // Returns an PjRtFuture<R> that is constructed from the result of invoking
  // functor `f` with *this value. If *this completes with an error, returned
  // future will also be an error.
  //
  // Sample usage: move-only type T passed by value
  //
  // std::move(future).Map<R>([](T value) -> U {
  //   return U(std::move(value)); // R must be constructible from U
  // })
  //
  template <typename R, typename F,
            typename U = std::invoke_result_t<
                F, std::conditional_t<is_move_only, T, const T&>>,
            std::enable_if_t<std::is_constructible_v<R, U>>* = nullptr>
  PjRtFuture<R> Map(F&& f) && {
    auto [promise, future] = PjRtFuture<R>::MakePromise();

    using Value = std::conditional_t<is_move_only, absl::StatusOr<T>,
                                     const absl::StatusOr<T>&>;
    std::move(*this).OnReady([promise = std::move(promise),
                              f = std::forward<F>(f)](Value value) mutable {
      if (ABSL_PREDICT_FALSE(promise.IsUniqueReference())) {
        promise.Set(Base::AbortedError());
      } else if (ABSL_PREDICT_TRUE(value.ok())) {
        if constexpr (is_move_only) {
          promise.emplace(absl::in_place_t{}, f(std::move(*value)));
        } else {
          promise.emplace(absl::in_place_t{}, f(*value));
        }
      } else {
        promise.Set(value.status());
      }
    });

    return std::move(future);
  }

  // Returns an PjRtFuture<R> that is constructed from the result of invoking
  // functor `f` with *this value. If *this completes with an error, returned
  // future will also be an error. Functor `f` must return a value of type
  // absl::StatusOr<U> where R is constructible from U. Returned absl::StatusOr
  // is automatically unwrapped and returned as a future payload.
  //
  // Sample usage:
  //
  // future.TryMap<R>([](const T& value) -> absl::StatusOr<U> {
  //   return U(value); // R must be constructible from U
  // })
  //
  template <
      typename R, typename F, typename U = std::invoke_result_t<F, const T&>,
      std::enable_if_t<!is_move_only && is_status_or<U> &&
                       std::is_constructible_v<R, typename U::value_type>>* =
          nullptr>
  PjRtFuture<R> TryMap(F&& f) const& {
    auto [promise, future] = PjRtFuture<R>::MakePromise();

    using Value = const absl::StatusOr<T>&;
    OnReady([promise = std::move(promise),
             f = std::forward<F>(f)](Value value) mutable {
      if (ABSL_PREDICT_FALSE(promise.IsUniqueReference())) {
        promise.Set(Base::AbortedError());
      } else if (ABSL_PREDICT_TRUE(value.ok())) {
        auto result = f(*value);
        if (ABSL_PREDICT_TRUE(result.ok())) {
          promise.emplace(absl::in_place_t{}, *std::move(result));
        } else {
          promise.Set(std::move(result).status());
        }
      } else {
        promise.Set(value.status());
      }
    });

    return std::move(future);
  }

  // Returns an PjRtFuture<R> that is constructed from the result of invoking
  // functor `f` with *this value. If *this completes with an error, returned
  // future will also be an error. Functor `f` must return a value of type
  // absl::StatusOr<U> where R is constructible from U. Returned absl::StatusOr
  // is automatically unwrapped and returned as a future payload.
  //
  // Sample usage: move-only type T passed by rvalue
  //
  // future.TryMap<R>([](const T& value) -> absl::StatusOr<U> {
  //   return U(value); // R must be constructible from U
  // })
  //
  template <typename R, typename F,
            typename U = std::invoke_result_t<
                F, std::conditional_t<is_move_only, T, const T&>>,
            std::enable_if_t<
                is_status_or<U> &&
                std::is_constructible_v<R, typename U::value_type>>* = nullptr>
  PjRtFuture<R> TryMap(F&& f) && {
    auto [promise, future] = PjRtFuture<R>::MakePromise();

    using Value = std::conditional_t<is_move_only, absl::StatusOr<T>,
                                     const absl::StatusOr<T>&>;
    std::move(*this).OnReady([promise = std::move(promise),
                              f = std::forward<F>(f)](Value value) mutable {
      if (ABSL_PREDICT_FALSE(promise.IsUniqueReference())) {
        promise.Set(Base::AbortedError());
      } else if (ABSL_PREDICT_TRUE(value.ok())) {
        auto result = [&] {
          if constexpr (is_move_only) {
            return f(std::move(*value));
          } else {
            return f(*value);
          }
        }();
        if (ABSL_PREDICT_TRUE(result.ok())) {
          promise.emplace(absl::in_place_t{}, *std::move(result));
        } else {
          promise.Set(std::move(result).status());
        }
      } else {
        promise.Set(value.status());
      }
    });

    return std::move(future);
  }

  // A `Map` overload that automatically infers the type of result from `f`.
  template <typename F, typename R = std::invoke_result_t<F, const T&>>
  PjRtFuture<R> Map(F&& f) const& {
    return Map<R>(std::forward<F>(f));
  }

  // A `Map` overload that automatically infers the type of result from `f`.
  template <typename F, typename R = std::invoke_result_t<
                            F, std::conditional_t<is_move_only, T, const T&>>>
  PjRtFuture<R> Map(F&& f) && {
    return std::move(*this).template Map<R>(std::forward<F>(f));
  }

  // A `TryMap` overload that automatically infers the type of result from `f`.
  template <typename F, typename R = std::invoke_result_t<F, const T&>,
            std::enable_if_t<is_status_or<R>>* = nullptr>
  PjRtFuture<typename R::value_type> TryMap(F&& f) const& {
    return TryMap<typename R::value_type>(std::forward<F>(f));
  }

  // A `TryMap` overload that automatically infers the type of result from `f`.
  template <typename F,
            typename R = std::invoke_result_t<
                F, std::conditional_t<is_move_only, T, const T&>>,
            std::enable_if_t<is_status_or<R>>* = nullptr>
  PjRtFuture<typename R::value_type> TryMap(F&& f) && {
    return std::move(*this).template TryMap<typename R::value_type>(
        std::forward<F>(f));
  }
};

// PjRtFuture<void> specialization for communicating stateless events.
//
// See PjRtFuture<T> documentation above for more details.
template <>
class PjRtFuture<void> : public internal::PjRtFutureBase<absl::Status> {
  using Base = internal::PjRtFutureBase<absl::Status>;

  template <typename U>
  static constexpr bool is_status_or =  // NOLINT
      tsl::internal::is_status_or_v<U>;

 public:
  class Promise : public Base::Promise {
   public:
    using Base::Promise::async_value;
    using Base::Promise::Promise;

    // Sets the promise completed with a given status. Must be called at most
    // once.
    //
    // After Set is called, completion event will be delivered to waiters on the
    // PjRtFuture constructed from a promise, via blocking or callbacks.
    void Set(absl::Status status = absl::OkStatus()) {
      Base::Promise::emplace(std::move(status));
    }

   private:
    friend class PjRtFuture<void>;
  };

  // This is a temporary class to support migration from CreatePromise() to
  // MakePromise() and an end goal of making Promise move-only type.
  class MoveOnlyPromise : public Promise {
   public:
    using Promise::Promise;
    using Promise::Set;

    MoveOnlyPromise(MoveOnlyPromise&&) = default;
    MoveOnlyPromise& operator=(MoveOnlyPromise&&) = default;

    // A helper function to convert move-only Promise to shared_ptr, which is
    // useful when the promise has to be captured by a std::function.
    std::shared_ptr<MoveOnlyPromise> ToShared() && {
      return std::make_shared<MoveOnlyPromise>(std::move(*this));
    }
  };

  // Returns a Promise that can be used to construct a PjRtFuture, and then Set
  // later.
  static Promise CreatePromise() {
    return Promise(tsl::MakeUnconstructedAsyncValueRef<absl::Status>());
  }

  // Returns a pair of connected Promise and PjRtFuture<>. Setting the returned
  // promise will fulfill the connected future.
  static std::pair<MoveOnlyPromise, PjRtFuture<>> MakePromise() {
    MoveOnlyPromise promise(
        tsl::MakeUnconstructedAsyncValueRef<absl::Status>());
    PjRtFuture<> future(promise);
    return std::make_pair(std::move(promise), std::move(future));
  }

  // Bring PjRtFutureBase constructors in scope.
  using Base::Base;

  // Constructor for unavailable future that will be fulfilled later via the
  // promise object.
  //
  // - on_block_start is called before Await starts to block.
  // - on_block_end is called after Await finishes blocking.
  explicit PjRtFuture(
      const Promise& promise,
      PjRtFutureHelpers::OnBlockStartFn on_block_start = nullptr,
      PjRtFutureHelpers::OnBlockEndFn on_block_end = nullptr)
      : Base(promise.ref(), std::move(on_block_start),
             std::move(on_block_end)) {}

  // Constructor for a future that is immediately ready with a given status.
  // For futures that are immediately ready with OK status, we use a global non
  // reference-counted async value that avoids heap allocation and reference
  // counting operations on a hot path.
  explicit PjRtFuture(absl::Status status)
      : Base(ABSL_PREDICT_TRUE(status.ok())
                 ? ready_promise_->AsRef()
                 : tsl::MakeAvailableAsyncValueRef<absl::Status>(
                       std::move(status)),
             /*on_block_start=*/nullptr, /*on_block_end=*/nullptr) {}

  using Base::Await;
  using Base::BlockUntilReady;
  using Base::OnReady;

  // Returns an PjRtFuture<R> that is constructed from the result of invoking
  // functor `f`. If *this completes with an error, returned future will also be
  // an error.
  //
  // Sample usage:
  //
  // future.Map<R>([]() -> U {
  //   return U(value); // R must be constructible from U
  // })
  //
  template <typename R, typename F, typename U = std::invoke_result_t<F>>
  PjRtFuture<R> Map(F&& f) {
    auto [promise, future] = PjRtFuture<R>::MakePromise();

    OnReady([promise = std::move(promise),
             f = std::forward<F>(f)](absl::Status status) mutable {
      if (ABSL_PREDICT_FALSE(promise.IsUniqueReference())) {
        promise.Set(Base::AbortedError());
      } else if (ABSL_PREDICT_TRUE(status.ok())) {
        promise.emplace(absl::in_place_t{}, f());
      } else {
        promise.Set(std::move(status));
      }
    });

    return std::move(future);
  }

  // Returns an PjRtFuture<R> that is constructed from the result of invoking
  // functor `f`. If *this completes with an error, returned future will also be
  // an error. Functor `f` must return a value of type absl::StatusOr<U> where R
  // is constructible from U. Returned absl::StatusOr is automatically unwrapped
  // and returned as a future payload.
  //
  // Sample usage:
  //
  // future.TryMap<R>([]() -> absl::StatusOr<U> {
  //   return U(value); // R must be constructible from U
  // })
  //
  template <typename R, typename F, typename U = std::invoke_result_t<F>,
            std::enable_if_t<
                is_status_or<U> &&
                std::is_constructible_v<R, typename U::value_type>>* = nullptr>
  PjRtFuture<R> TryMap(F&& f) {
    auto [promise, future] = PjRtFuture<R>::MakePromise();

    OnReady([promise = std::move(promise),
             f = std::forward<F>(f)](absl::Status status) mutable {
      if (ABSL_PREDICT_FALSE(promise.IsUniqueReference())) {
        promise.Set(Base::AbortedError());
      } else if (ABSL_PREDICT_TRUE(status.ok())) {
        auto result = f();
        if (ABSL_PREDICT_TRUE(result.ok())) {
          promise.emplace(absl::in_place_t{}, *std::move(result));
        } else {
          promise.Set(std::move(result).status());
        }
      } else {
        promise.Set(std::move(status));
      }
    });

    return std::move(future);
  }

  // A `Map` overload that automatically infers the type of result from `f`.
  template <typename F, typename R = std::invoke_result_t<F>>
  PjRtFuture<R> Map(F&& f) {
    return Map<R>(std::forward<F>(f));
  }

  // A `TryMap` overload that automatically infers the type of result from `f`.
  template <typename F, typename R = std::invoke_result_t<F>,
            std::enable_if_t<is_status_or<R>>* = nullptr>
  PjRtFuture<typename R::value_type> TryMap(F&& f) {
    return TryMap<typename R::value_type>(std::forward<F>(f));
  }

  // Returns an PjRtFuture<R> that is constructed from the given value. If *this
  // completes with an error, returned future will also be an error.
  //
  // Sample usage: make buffer available when future is ready
  //
  // std::unique_ptr<Buffer> buffer = ...;
  // future.MapTo<R>(std::move(buffer));
  template <typename R>
  PjRtFuture<absl::remove_cvref_t<R>> MapTo(R&& value) {
    return Map<absl::remove_cvref_t<R>>(
        [value = std::forward<R>(value)]() mutable {
          return std::move(value);
        });
  }

 private:
  // A promise that is immediately ready with OK status. Async value allocated
  // in the static storage and is not reference-counted.
  static absl::NoDestructor<tsl::AsyncValueOwningRef<absl::Status>>
      ready_promise_;
};

//===----------------------------------------------------------------------===//
// internal::PjRtFutureBase<T> implementation.
//===----------------------------------------------------------------------===//

namespace internal {

template <typename T, bool is_move_only>
PjRtFuture<> PjRtFutureBase<T, is_move_only>::GetReadyFuture() const {
  auto [promise, future] = PjRtFuture<>::MakePromise();
  promise_.AndThen(
      [self = promise_.AsPtr(), promise = std::move(promise)]() mutable {
        if constexpr (std::is_same_v<T, absl::Status>) {
          promise.Set(*self);
        } else {
          promise.Set(self->status());
        }
      });
  return std::move(future);
}

}  // namespace internal
}  // namespace xla

#endif  // XLA_PJRT_PJRT_FUTURE_H_
